!================================================================================!
! This file is part of crest.
!
! Copyright (C) 2018-2020 Stefan Grimme, Jakob Seibert
!
! crest is free software: you can redistribute it and/or modify it under
! the terms of the GNU Lesser General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! crest is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU Lesser General Public License for more details.
!
! You should have received a copy of the GNU Lesser General Public License
! along with crest.  If not, see <https://www.gnu.org/licenses/>.
!================================================================================!
subroutine zsort  
      use iomod
      use strucrd, only: rdnat,rdcoord,wrc0
      use miscdata, only: rcov
      implicit none
      integer n
      real*8,allocatable::xyz(:,:)                             ! comment out if subroutine
      real*8,allocatable::xyznew(:,:)
      real*8,allocatable :: tmp(:,:),tmpgeo(:,:)
      real*8,allocatable::geo(:,:)                             ! comment out if subroutine
      integer,allocatable::nat_mols(:),mvec(:),itmp(:),atnew(:)        
      integer,allocatable::na(:),nb(:),nc(:),at(:)             ! comment out if subroutine
      integer,allocatable::tmpna(:),tmpnb(:),tmpnc(:)          
      real*8,allocatable :: xx(:)     

      integer :: i,j,k,nn,nmol,refmol
      integer :: zmatcnt,zmatstart

      logical ex

      write(*,*)
      write(*,'(7x,''========================================'')')
      write(*,'(7x,''|             Z S O R T                |'')')
      write(*,'(7x,''|JS, Universitaet Bonn, MCTC, 05/2017  |'')')
      write(*,'(7x,''========================================'')')
      write(*,*)'input on file coord, Z matrix on file zmatrix'
      write(*,*)'sorted coord on file zcoord'

  allocate (xx(10))

  inquire (file='coord',exist=ex) ! comment out if subroutine
  if (.not. ex) stop 'ERROR: no coord file in folder!' ! comment out if subroutine

  call rdnat('coord',n)
  if (n < 4) then !without many atoms there is no point in sorting the zmat
    call copy('coord','zcoord')
    return
  end if

  allocate (xyz(3,n),at(n),mvec(n),xyznew(3,n),atnew(n), & ! comment out if subroutine
  &        geo(3,n),na(n),nb(n),nc(n))

  write (*,*) 'total number of atoms :',n ! comment out if subroutine

  call rdcoord('coord',n,at,xyz) ! comment out if subroutine

  call mrec(nmol,xyz,n,at,mvec)
  write (*,*) 'total number of frags :',nmol ! comment out if subroutine
  allocate (nat_mols(nmol))
  nat_mols = 0
  zmatcnt = 1
  zmatstart = 0
! check for largest fragment
  do i = 1,nmol
    nn = 0
    do j = 1,n
      if (mvec(j) .eq. i) nn = nn + 1
    end do
    nat_mols(i) = nn
  end do
  refmol = maxloc(nat_mols,1)
! loop over fragments
  do k = 1,nmol

    i = maxloc(nat_mols,1)
    write (*,*) "Fragment ",i ! comment out if subroutine
    allocate (itmp(nat_mols(i)),tmp(3,nat_mols(i)), &
&     tmpna(nat_mols(i)) &
&   ,tmpnb(nat_mols(i)),tmpnc(nat_mols(i)),tmpgeo(3,nat_mols(i)))
    nn = 0
    do j = 1,n
      if (mvec(j) .eq. i) then
        nn = nn + 1
        tmp(1:3,nn) = xyz(1:3,j)
        itmp(nn) = at(j)
      end if
    end do
    call recursive_zmat(tmp,itmp,nat_mols(i),tmpna,tmpnb,tmpnc,tmpgeo)
    call zmatpr(nat_mols(i),itmp,tmpgeo,tmpna,tmpnb,tmpnc,i) ! comment out if subroutine  ! write zmat file

    if (i .eq. refmol) then
      do j = 1,nat_mols(i)
        na(j) = tmpna(j)
        nb(j) = tmpnb(j)
        nc(j) = tmpnc(j)
        atnew(j) = itmp(j)
        xyznew(:,j) = tmp(:,j)
        zmatcnt = zmatcnt + 1
      end do
    else
      na(1 + zmatstart) = 1
      nb(1 + zmatstart) = 2
      nc(1 + zmatstart) = 3
      atnew(1 + zmatstart) = itmp(1)
      xyznew(:,1 + zmatstart) = tmp(:,1)
      zmatcnt = zmatcnt + 1
      if (nat_mols(i) .gt. 1) then
        na(2 + zmatstart) = tmpna(2) + zmatstart
        nb(2 + zmatstart) = 1
        nc(2 + zmatstart) = 2
        atnew(2 + zmatstart) = itmp(2)
        xyznew(:,2 + zmatstart) = tmp(:,2)
        zmatcnt = zmatcnt + 1
      end if
      if (nat_mols(i) .gt. 2) then
        na(3 + zmatstart) = tmpna(3) + zmatstart
        nb(3 + zmatstart) = tmpnb(3) + zmatstart
        nc(3 + zmatstart) = 1
        atnew(3 + zmatstart) = itmp(3)
        xyznew(:,3 + zmatstart) = tmp(:,3)
        zmatcnt = zmatcnt + 1
      end if
      do j = 4,nat_mols(i)
        na(j + zmatstart) = tmpna(j) + zmatstart
        nb(j + zmatstart) = tmpnb(j) + zmatstart
        nc(j + zmatstart) = tmpnc(j) + zmatstart
        atnew(j + zmatstart) = itmp(j)
        xyznew(:,j + zmatstart) = tmp(:,j)
        zmatcnt = zmatcnt + 1
      end do

    end if
    deallocate (itmp,tmp,tmpna,tmpnb,tmpnc,tmpgeo)
    zmatstart = zmatstart + nat_mols(i)
    nat_mols(i) = 0
  end do
  if (nmol .ne. 1) then
    !write(*,*) "Merged Z matrix"                           ! comment out if subroutine
    call xyzgeo(xyznew,n,na,nb,nc,1.0d0,geo)
    !call zmatpr(n,atnew,geo,na,nb,nc,0)                    ! comment out if subroutine
  end if

!     if you want to overwrite old xyz and at
!      xyz=xyznew
!      at=atnew

  call wrc0('zcoord',n,atnew,xyznew)

  write (*,*) 'terminated normally'

  deallocate (xyz,at,mvec,xyznew,atnew,geo,na,nb,nc,nat_mols)
!     deallocate(mvec,xyznew,atnew,nat_mols) if subroutine
  deallocate (xx)

end subroutine zsort

subroutine recursive_zmat(xyz,at,nat,na,nb,nc,geo)
  use crest_cn_module
  implicit none
!     recursive Z matrix generation
!     i/o block
  integer,intent(in) :: nat ! number of atoms
  integer,intent(in) :: at(nat) ! element number of atoms
  real*8,intent(in) :: xyz(3,nat) ! coordinates
  integer,intent(out) :: na(nat) ! column 1 of Z matrix (distance)
  integer,intent(out) :: nb(nat) ! column 2 of Z matrix (angle)
  integer,intent(out) :: nc(nat) ! column 3 of Z matrix (dihedral)
  real*8,intent(out) :: geo(3,nat) ! dist, ang and dihed in Z matrix
!     further variables
  real * 8 cn(nat),bond(nat,nat)
  integer i,zmatcnt,nref(nat),zmatpos(nat)
  logical taken(nat)
!     initialization
  geo = 0.0d0
  nref = 0
  taken = .false.
  zmatpos = 0
  zmatcnt = 1
  na = 0
  nb = 0
  nc = 0
!     code
  call calc_ncoord(nat,at,xyz,cn,bond)
  do i = 1,nat
!       j=maxloc(cn,1)
    if (.not. taken(i)) then

      taken(i) = .true.
      zmatpos(i) = zmatcnt
      call find_buddies(i,xyz,taken,nat,cn,bond,na,nb,nc, &
      &                 zmatcnt,nref,zmatpos)
    end if
  end do
  call zmatsort(nat,at,xyz,na,nb,nc,zmatpos) ! sort Z matrix
  call xyzgeo(xyz,nat,na,nb,nc,1.0d0,geo) ! compute dist, ang, dihed
end subroutine recursive_zmat

recursive subroutine find_buddies(i,xyz,taken,nat,cn,bond,na,nb,nc,zmatcnt,nref,zmatpos)
  implicit none
  real * 8 xyz(3,nat),cn(nat),bond(nat,nat)
  integer i,nat,j,icn,k,na(nat),nb(nat),nc(nat)
  integer zmatcnt,nref(nat),zmatpos(nat)
  logical taken(nat)

  icn = nint(cn(i))

  do k = 1,icn
    j = maxloc(bond(:,i),1)
    bond(j,i) = 0.0d0

    if (i .eq. j) cycle
    if (.not. taken(j)) then
      taken(j) = .true.
      zmatcnt = zmatcnt + 1
      if (zmatcnt .gt. nat) exit
      zmatpos(j) = zmatcnt
      if (nref(i) .ne. 0) then
        if (nref(nref(i)) .eq. 0) then
          na(j) = na(nref(i))
          if (zmatcnt .le. 3 .or. nb(nref(i)) .eq. 0) then !
            nb(j) = nref(i) !
            nref(i) = j !
          else !
            nb(j) = nb(nref(i))
          end if !
          if (zmatcnt .le. 3) then !
            nc(j) = 0 !
          else !
            nc(j) = nref(i)
          end if
        else if (nref(nref(i)) .ne. 0) then
          na(j) = i
          nb(j) = nref(i)
          if (nb(nref(i)) .eq. 0) then
            nc(j) = nref(nref(i))
          else
            nc(j) = nb(nref(i))
          end if
          nref(i) = j
        else
          na(j) = i
          nb(j) = na(i)
          nc(j) = nb(i)
          nref(i) = j
        end if
      else
        na(j) = i
        if (na(i) .eq. 0 .and. nref(i) .ne. 0) then !
          nb(j) = nref(i) !
        else !
          nb(j) = na(i)
        end if !
        nc(j) = nb(i)
        nref(i) = j
      end if
      call find_buddies(j,xyz,taken,nat,cn,bond,na,nb,nc,zmatcnt,nref,zmatpos)
    end if
  end do
end subroutine find_buddies

subroutine zmatsort(nat,at,xyz,na,nb,nc,zmatpos)
  implicit none
!     sort zmat according to zmat position from previous subroutine
  integer nat,at(nat),na(nat),nb(nat),nc(nat),zmatpos(nat)
  integer tmpat(nat),tmpna(nat),tmpnb(nat),tmpnc(nat)
  integer i,oldpos
  real * 8 xyz(3,nat),tmpxyz(3,nat)
  tmpat = 0
  tmpna = 0
  tmpnb = 0
  tmpnc = 0
  tmpxyz = 0.0d0
  do i = nat,1,-1

    oldpos = maxloc(zmatpos,1)
    tmpat(i) = at(oldpos)
    if (na(oldpos) .ne. 0) then
      tmpna(i) = zmatpos(na(oldpos))
    else
      tmpna(i) = 0
    end if
    if (nb(oldpos) .ne. 0) then
      tmpnb(i) = zmatpos(nb(oldpos))
    else
      tmpnb(i) = 0
    end if
    if (nc(oldpos) .ne. 0) then
      tmpnc(i) = zmatpos(nc(oldpos))
    else
      tmpnc(i) = 0
    end if

    tmpxyz(:,i) = xyz(:,oldpos)

    zmatpos(oldpos) = 0

  end do

  at = tmpat
  na = tmpna
  nb = tmpnb
  nc = tmpnc
  xyz = tmpxyz
!     check if Z matrix reasonable
  i = 1
  if (i .eq. 1) then
    if (na(1) .ne. 0 .or. nb(1) .ne. 0  &
&      .or. nc(1) .ne. 0) stop "ERROR: first Z matrix  &
&     entry is erroneous"
    if (nat .gt. 1) then
      if (na(2) .eq. 0 .or. nb(2) .ne. 0  &
&         .or. nc(2) .ne. 0) stop "ERROR: second Z matrix &
&         entry is erroneous"
      if (nat .gt. 2) then
        if (na(3) .eq. 0 .or. nb(3) .eq. 0  &
&            .or. nc(3) .ne. 0) stop "ERROR: third Z matrix &
&            entry is erroneous"
        if (nat .gt. 3) then
          do i = 4,nat
            if (na(i) .eq. 0 .or. nb(i) .eq. 0 .or. nc(i) .eq. 0) then
              write (*,*) "ERROR: Erroneous entry in Z matrix at position ",i
              stop
            end if

          end do
        end if
      end if
    end if
  end if
end subroutine zmatsort

!------------------------------------------------------------------------
subroutine zmatpr(nat,at,geo,na,nb,nc,molnum)
  use crest_parameters, only: wp, bohr
  use strucrd,only:i2e
  implicit none
  integer :: nat,na(nat),nb(nat),nc(nat),at(nat)
  real(wp) :: geo(3,nat)
  character(len=20) :: filename
  integer :: i,l,m,n,molnum
  real(wp) :: bl,ang,dihed,pi
  parameter(pi=3.14159265358979D0)
  do i = 1,nat
    l = 1
    m = 1
    n = 1
    if (i .eq. 1) then
      l = 0
      m = 0
      n = 0
    end if
    if (i .eq. 2) then
      m = 0
      n = 0
    end if
    if (i .eq. 3) n = 0
    bl = geo(1,i) * bohr
    ang = geo(2,i) * 180./pi
    dihed = geo(3,i) * 180./pi
    if (dihed .gt. 180.0d0) dihed = dihed - 360.0d0
    write (*,'(i4,2x,a2,f12.6,2x,f10.4,2x,f10.4,i6,2i5)') &
    &   i,i2e(at(i)),bl,ang,dihed,na(i),nb(i),nc(i)
  end do

  write (filename,'("zmatrix",i0,".zmat")') molnum
  open (unit=42,file=filename)

  write (42,'(a2)') i2e(at(1))
  if (nat .gt. 1) then
    write (42,'(a2,x,i0,x,f8.3)') i2e(at(2)),na(2),geo(1,2) * bohr
  end if
  if (nat .gt. 2) then
    write (42,'(a2,x,i0,x,f8.3,x,i0,x,f8.3)') i2e(at(3)),na(3) &
    &                  ,geo(1,3) * bohr,nb(3),geo(2,3) * 180./pi
  end if

  do i = 4,nat
    bl = geo(1,i) * bohr
    ang = geo(2,i) * 180./pi
    dihed = geo(3,i) * 180./pi
    if (dihed .gt. 180.0d0) dihed = dihed - 360.0d0
    write (42,'(a2,x,i0,x,f8.3,x,i0,x,f8.3,x,i0,x,f8.3)') &
    &         i2e(at(i)),na(i),bl,nb(i),ang,nc(i),dihed
  end do
  write (42,*)
  close (42)
end subroutine zmatpr
